#include "TCanvas.h"
#include "TH2F.h"
#include "TTree.h"
#include "TFile.h"
#include "TStyle.h"
#include "TColor.h"
#include "TF2.h"
#include "TExec.h"
#include "TCanvas.h"
#include "TSystem.h"
#include "../../MarvelSim/source/AnalysisManager/include//SimEvent.h"
const int el = 5;
const TString tail = "MyDNA";

void applyYOLO()
{
    gSystem->Load("../../MarvelSim/build/libSimEventDict.so");
    TString str_el[6] = {"H_Z1A1", "He_Z2A4", "C_Z6A12", "O_Z8A16", "D_Z1A2", "CheckTails"};
    //TFile* mc_file = new TFile(Form("~/WorkSpace/Marvel/NRBKG/%s_wide/%s_feature.root", str_el[el].Data(), str_el[el].Data()));
    //TFile* mc_file = new TFile(Form("/home/mart/WorkSpace/Marvel/NRBKG/CheckTails/H_%s_feature.root", tail.Data()));
    //TTree* mc_tree = (TTree*)mc_file->Get("t_data");
    TChain* mc_tree = new TChain("t_data");
    mc_tree->Add("/home/mart/WorkSpace/Marvel/NRBKG/CheckTails/H_10keV_feature.root");
    mc_tree->Add("/home/mart/WorkSpace/Marvel/NRBKG/CheckTails/H_10keV_2_feature.root");
    mc_tree->Add("/home/mart/WorkSpace/Marvel/NRBKG/CheckTails/H_10keV_3_feature.root");
    mc_tree->Add("/home/mart/WorkSpace/Marvel/NRBKG/CheckTails/H_10keV_4_feature.root");
    mc_tree->Add("/home/mart/WorkSpace/Marvel/NRBKG/CheckTails/H_10keV_5_feature.root");
    TFile* yolo_file = new TFile(Form("%s_wide_yolo.root", str_el[el].Data()));
    TTree* yolo_tree = (TTree*)yolo_file->Get("yolo_output");

    TFile* new_file = new TFile(Form("%s_wide_WithYOLO.root", str_el[el].Data()), "RECREATE");
    TTree* new_tree = mc_tree->CloneTree(0);

    Int_t EntryNum = 0;
    Int_t ERNum = 0;
    Int_t NRNum = 0;
    Int_t CosmicNum = 0;

    Long64_t yolo_EntryNum = 0;
    Long64_t yolo_ERNum = 0;
    Long64_t yolo_NRNum = 0;
    Long64_t yolo_CosmicNum = 0;

    yolo_tree->SetBranchAddress("EntryNum", &yolo_EntryNum);
    yolo_tree->SetBranchAddress("ERNum", &yolo_ERNum);
    yolo_tree->SetBranchAddress("NRNum", &yolo_NRNum);
    yolo_tree->SetBranchAddress("CosmicNum", &yolo_CosmicNum);

    new_tree->Branch("ERNum", &ERNum, "ERNum/I");
    new_tree->Branch("NRNum", &NRNum, "NRNum/I");
    new_tree->Branch("CosmicNum", &CosmicNum, "CosmicNum/I");

    Long64_t num_of_yolo = yolo_tree->GetEntries();

    for(Long64_t iy = 0; iy < num_of_yolo; ++iy)
    {
        yolo_tree->GetEntry(iy);
        EntryNum = yolo_EntryNum;
        ERNum = yolo_ERNum;
        NRNum = yolo_NRNum;
        CosmicNum = yolo_CosmicNum;
        mc_tree->GetEntry(EntryNum);
        new_tree->Fill();
    }

    new_file->cd();
    new_tree->AutoSave();
    new_file->Close();
}
